%% Author: zhangjun
%% Created: 2009-3-29
%% Description: TODO: Add description to pickle
-module(pickle).

%%
%% Include files
%%

%%
%% Exported Functions
%%
-export([test/0, unpickle/2, pickle/2]).

%%
%% API Functions
%%

%%
%% TODO: Add description of test/function_arity
%%
test() -> 
	ok.
%%
%% TODO: Add description of unpickle/function_arity
%% first parameter is byte/0 or short/0 or list/2 and so on...
unpickle({_, Pickler}, Bin) ->
	element(1, Pickler(Bin)).
%%
%% TODO: Add description of pickle/function_arity
%%
pickle({Pickler, _}, Value) ->
	lists:reverse(Pickler([], Value)).

%%
%% Local Functions
%%

%%% Byte

byte() -> 
	{fun write_byte/2, fun read_byte/1}.

write_byte(Acc, Byte) -> 
	[<<Byte:8>>|Acc].

read_byte(Bin) -> 
	<<Byte:8, Rest/binary>> = Bin,
	{Byte, Rest}.

%%% Unsigned short

short() -> 
	{fun write_short/2, fun read_short/1}.

write_short(Acc, Word) -> 
	[<<Word:16/little>>|Acc].

read_short(Bin) -> 
	<<Word:16/little, Rest/binary>> = Bin,
	{Word, Rest}.

%%% Signed short

sshort() -> 
	{fun write_sshort/2, fun read_sshort/1}.

write_sshort(Acc, Word) -> 
	[<<Word:16/little-signed>>|Acc].

read_sshort(Bin) -> 
	<<Word:16/little-signed, Rest/binary>> = Bin,
	{Word, Rest}.

%%% List. We supply a pickler for list length 
%%% as well as a pickler for list elements.
%%% Elem is byte or short or sshort and so on...
%%% e.g. elem is byte: while write_list Elem is write_byte, read_list Elem is read_byte
list(Len, Elem) ->
	{fun(Acc, List) -> write_list(Len, Elem, Acc, List) end, 
	 fun(Bin) -> read_list(Len, Elem, Bin) end }.

write_list({Len, _}, {Elem, _}, Acc, List) ->
	Acc1 = Len(Acc, length(List)),
	Fun = fun(A, Acc2) -> Elem(Acc2, A) end,
	lists:foldr(Fun, Acc1, List).

read_list({_, Len}, {_, Elem}, Bin) ->
	{N, Bin1} = Len(Bin),
	read_list(N, [], Elem, Bin1).

read_list(0, Acc, _, Bin) -> {Acc, Bin};
read_list(N, Acc, Elem, Bin) ->
	{E, Bin1} = Elem(Bin),
	read_list(N - 1, [E|Acc], Elem, Bin1).

%%% Alternative selection. This could probably use some
%%% deeper thinking. Otherwise, we take a pickler for the tag
%%% as well as a tuple of two functions. The first one
%%% returns the tag value and a pickler based on the supplied
%%% value. The second one selects a pickler based on a tag value.

choice(Tag, Choice) ->
	{fun(Acc, Value) -> write_choice(Tag, Choice, Acc, Value) end,
	 fun(Bin) -> read_choice(Tag, Choice, Bin) end }.

write_choice({Tag, _}, {Choice, _}, Acc, Value) 
  when is_function(Tag), 
	   is_function(Choice) ->
	{T, {Pickler, _}} = Choice(Value),
	Acc1 = Tag(Acc, T),
	Pickler(Acc1, Value).

read_choice({_, Tag}, {_, Choice}, Bin) 
  when is_function(Tag), 
	   is_function(Choice) ->
	{T, Bin1} = Tag(Bin),
	{_, Pickler} = Choice(T),
	Pickler(Bin1).

%%% Optional value. Use 'none' to indicate no value.

optional(Pickler) ->
	{fun(Acc, Value) -> write_optional(Pickler, Acc, Value) end,
	 fun(Bin) -> read_optional(Pickler, Bin) end}.

write_optional(_, Acc, none) ->
	[<<0>>|Acc];

write_optional({Pickler, _}, Acc, Value) ->
	Pickler([<<1>>|Acc], Value).

read_optional({_, Pickler}, Bin) ->
	<<Opt:8, Bin1/binary>> = Bin,
	case Opt of 
		0 -> {none, Bin1};
		_ -> Pickler(Bin1)
	end.

%%% Wrapper. Take a pickler and a wrapper tuple of two functions
%%% where the first one is used to convert the value before 
%%% pickling and the second one after unpickling.

wrap(Wrap, Pickler) ->
	{fun(Acc, Value) -> write_wrap(Wrap, Pickler, Acc, Value) end,
	 fun(Bin) -> read_wrap(Wrap, Pickler, Bin) end}.

write_wrap({Wrap, _}, {Pickler, _}, Acc, Value) ->
	Pickler(Acc, Wrap(Value)).

read_wrap({_, Wrap}, {_, Pickler}, Bin) ->
	{Value, Bin1} = Pickler(Bin),
	{Wrap(Value), Bin1}.

%%%elem enumeration

enum(Enum, Pickler) ->
	wrap(wrap_enum(Enum), Pickler).

wrap_enum(Enum) 
  when is_tuple(Enum) ->
	wrap_enum_1(prep_enum_tuple(Enum));

wrap_enum(Enum) 
  when is_list(Enum) ->
	wrap_enum_1(prep_enum_list(Enum)).

%% convert {cow, sheep, horse} into a pair (tuple of two elements) of lists 
%%[{cow, 1}, {sheep, 2}, {horse, 3}] and [{1, cow}, {2, sheep}, {3, horse}]. 

prep_enum_tuple(Enum)
  when is_tuple(Enum) ->
	prep_enum_tuple(Enum, size(Enum), [], []).

prep_enum_tuple(_, 0, Acc1, Acc2) ->
    {Acc1, Acc2};

prep_enum_tuple(Enum, N, Acc1, Acc2) ->
	prep_enum_tuple(Enum, N - 1, 
					[{element(N, Enum), N}|Acc1],
					[{N, element(N, Enum)}|Acc2]).

prep_enum_list(Enum) 
  when is_list(Enum) ->
	% expect a list of {tag, #value} pairs
	Inv = fun({Key, Val}) -> {Val, Key} end,
	InvEnum = lists:map(Inv, Enum),
	{Enum, InvEnum}.

wrap_enum_1({List1, List2}) ->
	F = fun(A, B) -> A < B end,
	%% gb_trees needs an ordered list
	Dict1 = lists:sort(F, List1),
	Dict2 = lists:sort(F, List2),
	Tree1 = gb_trees:from_orddict(Dict1),
	Tree2 = gb_trees:from_orddict(Dict2),
	{fun(Key) -> gb_trees:get(Key, Tree1) end,
	 fun(Key) -> gb_trees:get(Key, Tree2) end}.

%%% Tuple. Uses a tuple of picklers of the same size.

tuple(Picklers) 
  when is_tuple(Picklers) ->
	wrap({fun tuple_to_list/1, 
		  fun list_to_tuple/1}, 
		 tuple_0(tuple_to_list(Picklers))).

record(Tag, Picklers) 
  when is_tuple(Picklers) ->
	wrap({fun(Record) -> record_to_list(Tag, Record) end,
		  fun(List) -> list_to_record(Tag, List) end}, 
		 tuple_0(tuple_to_list(Picklers))).

write_tuple_0([], Acc, _) ->
	Acc;

%%% every tuple element corrspond to a Pickler 
write_tuple_0([{Pickler, _}|Rest], Acc, [Value|Tuple]) ->
	write_tuple_0(Rest, Pickler(Acc, Value), Tuple).

read_tuple_0(Picklers, Bin) ->
	read_tuple_0(Picklers, Bin, []).

read_tuple_0([], Bin, Acc) ->
	{lists:reverse(Acc), Bin};

read_tuple_0([{_, Pickler}|Rest], Bin, Acc) ->
	{Value, Bin1} = Pickler(Bin),
	read_tuple_0(Rest, Bin1, [Value|Acc]).

%%% It's convenient to be able to convert the tuple
%%% to a list first as there's no erlang:prepend_element/2.

tuple_0(Picklers) 
  when is_list(Picklers) ->
	{fun(Acc, Value) -> write_tuple_0(Picklers, Acc, Value) end,
	 fun(Bin) -> read_tuple_0(Picklers, Bin) end}.

record_to_list(Tag, Record) 
  when is_atom(Tag) ->
	lists:nthtail(1, tuple_to_list(Record)).

list_to_record(Tag, List) 
  when is_atom(Tag), 
	   is_list(List) ->
	list_to_tuple([Tag|List]).

